package com.interview.rtc.aigc.utils;

import lombok.Data;
import org.apache.commons.codec.binary.Hex;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.text.SimpleDateFormat;
import java.time.ZoneId;
import java.time.ZonedDateTime;
import java.time.format.DateTimeFormatter;
import java.util.*;

/**
 * 火山引擎API签名工具
 */
public class VolcSigner {

    private static final String CONTENT_TYPE = "application/json";
    private static final String HOST = "rtc.volcengineapi.com";
    private static final String SERVICE = "rtc";
    private static final String REGION = "cn-north-1";
    private static final String SCHEMA = "https";
    private static final String PATH = "/";

    private static final BitSet URLENCODER = new BitSet(256);
    private static final String CONST_ENCODE = "0123456789ABCDEF";

    static {
        int i;
        for (i = 97; i <= 122; ++i) {
            URLENCODER.set(i);
        }

        for (i = 65; i <= 90; ++i) {
            URLENCODER.set(i);
        }

        for (i = 48; i <= 57; ++i) {
            URLENCODER.set(i);
        }
        URLENCODER.set('-');
        URLENCODER.set('_');
        URLENCODER.set('.');
        URLENCODER.set('~');
    }

    @Data
    public static class VolcCredentials {
        private String accessKeyId;
        private String secretKey;
        private String sessionToken;
    }

    /**
     * 生成火山引擎API请求的签名信息
     *
     * @param credentials 凭证信息
     * @param action API动作名称
     * @param version API版本
     * @param requestBody 请求体JSON
     * @return 包含签名信息的请求头Map
     */
    public static Map<String, String> generateSignedHeaders(VolcCredentials credentials,
                                                          String action,
                                                          String version,
                                                          String requestBody) {
        try {
            // 获取当前时间
            ZonedDateTime now = ZonedDateTime.now(ZoneId.of("UTC"));
            String xDate = now.format(DateTimeFormatter.ofPattern("yyyyMMdd'T'HHmmss'Z'"));
            String shortXDate = xDate.substring(0, 8);

            // 计算请求体哈希
            String xContentSha256 = hashSHA256(requestBody.getBytes(StandardCharsets.UTF_8));

            // 构建请求头
            Map<String, String> headers = new TreeMap<>();
            headers.put("Content-Type", CONTENT_TYPE);
            headers.put("Host", HOST);
            headers.put("X-Date", xDate);
            headers.put("X-Content-Sha256", xContentSha256);

            if (credentials.getSessionToken() != null && !credentials.getSessionToken().isEmpty()) {
                headers.put("X-Security-Token", credentials.getSessionToken());
            }

            // 构建查询参数
            Map<String, String> queryParams = new TreeMap<>();
            queryParams.put("Action", action);
            queryParams.put("Version", version);

            // 构建规范请求
            String canonicalRequest = buildCanonicalRequest("POST", queryParams, headers, xContentSha256);

            // 计算签名
            String credentialScope = shortXDate + "/" + REGION + "/" + SERVICE + "/request";
            String stringToSign = buildStringToSign(xDate, credentialScope, canonicalRequest);
            byte[] signKey = genSigningSecretKeyV4(credentials.getSecretKey(), shortXDate, REGION, SERVICE);
            String signature = Hex.encodeHexString(hmacSHA256(signKey, stringToSign));

            // 构建Authorization头
            String signHeader = "host;x-date;x-content-sha256;content-type";
            String authorization = String.format("HMAC-SHA256 Credential=%s/%s, SignedHeaders=%s, Signature=%s",
                    credentials.getAccessKeyId(), credentialScope, signHeader, signature);

            // 添加到请求头
            headers.put("Authorization", authorization);

            return headers;
        } catch (Exception e) {
            throw new RuntimeException("Generate signed headers failed", e);
        }
    }

    private static String buildCanonicalRequest(String method, Map<String, String> queryParams,
            Map<String, String> headers, String xContentSha256) throws UnsupportedEncodingException {
        StringBuilder canonicalRequest = new StringBuilder();

        // 1. HTTP方法
        canonicalRequest.append(method).append("\n");

        // 2. 规范URI
        canonicalRequest.append(PATH).append("\n");

        // 3. 规范查询字符串
        StringBuilder querySB = new StringBuilder();
        for (String key : queryParams.keySet()) {
            querySB.append(signStringEncoder(key))
                  .append("=")
                  .append(signStringEncoder(queryParams.get(key)))
                  .append("&");
        }
        if (querySB.length() > 0) {
            querySB.deleteCharAt(querySB.length() - 1);
        }
        canonicalRequest.append(querySB).append("\n");

        // 4. 规范请求头
        for (Map.Entry<String, String> entry : headers.entrySet()) {
            canonicalRequest.append(entry.getKey().toLowerCase())
                          .append(":")
                          .append(entry.getValue().trim())
                          .append("\n");
        }
        canonicalRequest.append("\n");

        // 5. 已签名的请求头
        canonicalRequest.append("host;x-date;x-content-sha256;content-type").append("\n");

        // 6. 请求体哈希
        canonicalRequest.append(xContentSha256);

        return canonicalRequest.toString();
    }

    private static String buildStringToSign(String xDate, String credentialScope, String canonicalRequest)
            throws Exception {
        String hashedRequest = hashSHA256(canonicalRequest.getBytes(StandardCharsets.UTF_8));
        return String.format("HMAC-SHA256\n%s\n%s\n%s", xDate, credentialScope, hashedRequest);
    }

    private static String signStringEncoder(String source) {
        if (source == null) {
            return null;
        }
        StringBuilder buf = new StringBuilder(source.length());
        ByteBuffer bb = StandardCharsets.UTF_8.encode(source);
        while (bb.hasRemaining()) {
            int b = bb.get() & 255;
            if (URLENCODER.get(b)) {
                buf.append((char) b);
            } else if (b == 32) {
                buf.append("%20");
            } else {
                buf.append("%");
                char hex1 = CONST_ENCODE.charAt(b >> 4);
                char hex2 = CONST_ENCODE.charAt(b & 15);
                buf.append(hex1);
                buf.append(hex2);
            }
        }
        return buf.toString();
    }

    private static String hashSHA256(byte[] content) throws Exception {
        try {
            MessageDigest md = MessageDigest.getInstance("SHA-256");
            return Hex.encodeHexString(md.digest(content));
        } catch (Exception e) {
            throw new Exception("Unable to compute hash while signing request: " + e.getMessage(), e);
        }
    }

    private static byte[] hmacSHA256(byte[] key, String content) throws Exception {
        try {
            Mac mac = Mac.getInstance("HmacSHA256");
            mac.init(new SecretKeySpec(key, "HmacSHA256"));
            return mac.doFinal(content.getBytes(StandardCharsets.UTF_8));
        } catch (Exception e) {
            throw new Exception("Unable to calculate a request signature: " + e.getMessage(), e);
        }
    }

    private static byte[] genSigningSecretKeyV4(String secretKey, String date, String region, String service)
            throws Exception {
        byte[] kDate = hmacSHA256(("HMAC" + secretKey).getBytes(StandardCharsets.UTF_8), date);
        byte[] kRegion = hmacSHA256(kDate, region);
        byte[] kService = hmacSHA256(kRegion, service);
        return hmacSHA256(kService, "request");
    }
}
